// Problem 196: Prime triplets

package main

import (
	"fmt"
	"math/big"
)

func firstElementOfRow(n int64) int64 {
	return n*(n-1)/2 + 1
}

func isTriplet(n, r int64) bool {
	// only for odd row
	var cell = [9]int64{n - r + 1, n - 2*r + 2, n - 2*r + 4, n + r - 1, n - 2, n + 2*r, n + r + 1, n + 2*r + 2, n + 2}
	// only for even row
	//var cell = [9]int64{n - r, n - 2*r + 2, n - 2, n - r + 2, n - 2*r + 4, n + 2, n + r, n + 2*r, n + 2*r + 2}
	bn := big.NewInt(int64(n))

	if !bn.ProbablyPrime(2) {
		return false
	}
	count := 0
	for t := 0; t < 3; t++ {
		bn.SetInt64(cell[t*3])
		if bn.ProbablyPrime(2) {
			count++
			bn.SetInt64(cell[t*3+1])
			if bn.ProbablyPrime(2) {
				return true
			}
			bn.SetInt64(cell[t*3+2])
			if bn.ProbablyPrime(2) {
				return true
			}
		}
	}
	return count >= 2
}

func p196() {
	var lines []int64 = []int64{5678027, 7208785}
	var sum int64 = 0
	for _, line := range lines {
		start := firstElementOfRow(line)
		end := start + line - 1
		if start%2 == 0 {
			start = start + 1
		}
		for n := start; n <= end; n += 2 {
			if isTriplet(n, line) {
				sum += n
			}
		}
	}
	fmt.Println("Problem 196:", sum)
}
